from PySide import QtCore, QtGui
import diabind

# ------------------------------------------------------------------------------

class BuildFieldSimple(object):
    def __init__(self, name):
        self.name = name

    def getValue(self, obj):
        return getattr(obj, "get_" + self.name)()

    def getValuleString(self, val):
        return str( val )

    def addValueData(self, item, val):
        item.setText(1, self.getValuleString(val))

    def createTreeItem(self, treeItem, val):
        return QtGui.QTreeWidgetItem(treeItem)

    def do(self, obj, treeItem):
        val = None
        try:
            val = self.getValue(obj)
        except diabind.DiaError:
            return

        item = self.createTreeItem(treeItem, val)
        item.setText(0, self.name)
        self.addValueData(item, val)

class BuildFieldHex(BuildFieldSimple):
    def getValuleString(self, val):
        return hex( val )

class BuildFieldRva(BuildFieldHex):
    def createTreeItem(self, treeItem, val):
        return TreeItemSymRva(treeItem, val)

class BuildFieldEnum(BuildFieldSimple):
    def __init__(self, name, enumName):
        super(BuildFieldEnum, self).__init__(name)
        self.enumName = enumName

    def getValuleString(self, val):
        try:
            return str( getattr(diabind, self.enumName).values[val] )
        except KeyError:
            return self.enumName + "(" + str(val) + ")"

class BuildFieldUndecName(BuildFieldSimple):
    def getValue(self, obj):
        return getattr(obj, "get_" + self.name)(diabind.UndecorateFlag_Complete)

class BuildFieldSymbol(BuildFieldSimple):
    def addValueData(self, item, val):
        TreeItemSym(item, val)

class BuildFieldSrcFile(BuildFieldSimple):
    def addValueData(self, item, val):
        TreeItemSrcFile(item, val)

class BuildFieldEnumSymbols(BuildFieldSimple):
    def addValueData(self, item, val):
        for sym in val:
            TreeItemSym(item, sym)

# ------------------------------------------------------------------------------

_SYMBOL_FIELD_BUILDER   = ( BuildFieldEnum( "symTag", "SymTag" ),
                            BuildFieldSimple( "name" ),
                            BuildFieldUndecName( "undecoratedNameEx" ),
                            BuildFieldSimple( "symIndexId" ),
                            BuildFieldSymbol( "lexicalParent" ),
                            BuildFieldSymbol( "classParent" ),
                            BuildFieldSymbol( "type" ),
                            BuildFieldEnum( "dataKind", "DataKind" ),
                            BuildFieldEnum( "locationType", "LocationType" ),
                            BuildFieldSimple( "addressSection" ),
                            BuildFieldHex( "addressOffset" ),
                            BuildFieldRva( "relativeVirtualAddress" ),
                            BuildFieldEnum( "registerId", "CV_HREG_e" ),
                            BuildFieldHex( "offset" ),
                            BuildFieldSimple( "length" ),
                            BuildFieldSimple( "slot" ),
                            BuildFieldSimple( "volatileType" ),
                            BuildFieldSimple( "constType" ),
                            BuildFieldSimple( "unalignedType" ),
                            BuildFieldEnum( "access", "CV_access_e" ),
                            BuildFieldSimple( "libraryName" ),
                            BuildFieldEnum( "platform", "CV_CPU_TYPE_e" ),
                            BuildFieldEnum( "language", "CV_CFL_LANG" ),
                            BuildFieldSimple( "editAndContinueEnabled" ),
                            BuildFieldSimple( "frontEndMajor" ),
                            BuildFieldSimple( "frontEndMinor" ),
                            BuildFieldSimple( "frontEndBuild" ),
                            BuildFieldSimple( "backEndMajor" ),
                            BuildFieldSimple( "backEndMinor" ),
                            BuildFieldSimple( "backEndBuild" ),
                            BuildFieldSimple( "sourceFileName" ),
                            BuildFieldEnum( "thunkOrdinal", "THUNK_ORDINAL" ),
                            BuildFieldSimple( "thisAdjust" ),
                            BuildFieldHex( "virtualBaseOffset" ),
                            BuildFieldSimple( "virtual" ),
                            BuildFieldSimple( "intro" ),
                            BuildFieldSimple( "pure" ),
                            BuildFieldSimple( "value" ),
                            BuildFieldEnum( "callingConvention", "CV_call_e" ),
                            BuildFieldEnum( "baseType", "BasicType" ),
                            BuildFieldSimple( "token" ),
                            BuildFieldHex( "timeStamp" ),
                            BuildFieldSimple( "guid" ),
                            BuildFieldSimple( "symbolsFileName" ),
                            BuildFieldSimple( "reference" ),
                            BuildFieldSimple( "count" ),
                            BuildFieldSimple( "bitPosition" ),
                            BuildFieldSymbol( "arrayIndexType" ),
                            BuildFieldSimple( "packed" ),
                            BuildFieldSimple( "constructor" ),
                            BuildFieldSimple( "overloadedOperator" ),
                            BuildFieldSimple( "nested" ),
                            BuildFieldSimple( "hasNestedTypes" ),
                            BuildFieldSimple( "hasAssignmentOperator" ),
                            BuildFieldSimple( "hasCastOperator" ),
                            BuildFieldSimple( "scoped" ),
                            BuildFieldSimple( "virtualBaseClass" ),
                            BuildFieldSimple( "indirectVirtualBaseClass" ),
                            BuildFieldHex( "virtualBasePointerOffset" ),
                            BuildFieldSymbol( "virtualTableShape" ),
                            BuildFieldSimple( "code" ),
                            BuildFieldSimple( "function" ),
                            BuildFieldSimple( "managed" ),
                            BuildFieldSimple( "msil" ),
                            BuildFieldSimple( "virtualBaseDispIndex" ),
                            BuildFieldHex( "age" ),
                            BuildFieldHex( "signature" ),
                            BuildFieldSimple( "compilerGenerated" ),
                            BuildFieldSimple( "addressTaken" ),
                            BuildFieldSimple( "rank" ),
                            BuildFieldSimple( "lowerBound" ),
                            BuildFieldSimple( "upperBound" ),
                            BuildFieldSimple( "dataBytes" ),
                            BuildFieldSimple( "targetSection" ),
                            BuildFieldHex( "targetOffset" ),
                            BuildFieldHex( "targetRelativeVirtualAddress" ),
                            BuildFieldEnum( "machineType", "IMAGE_FILE_MACHINE" ),
                            BuildFieldSimple( "oemId" ),
                            BuildFieldSimple( "oemSymbolId" ),
                            BuildFieldSimple( "types" ),
                            BuildFieldSymbol( "objectPointerType" ),
                            BuildFieldEnum( "udtKind", "UdtKind" ),
                            BuildFieldSimple( "noReturn" ),
                            BuildFieldSimple( "customCallingConvention" ),
                            BuildFieldSimple( "noInline" ),
                            BuildFieldSimple( "optimizedCodeDebugInfo" ),
                            BuildFieldSimple( "notReached" ),
                            BuildFieldSimple( "interruptReturn" ),
                            BuildFieldSimple( "farReturn" ),
                            BuildFieldSimple( "isStatic" ),
                            BuildFieldSimple( "hasDebugInfo" ),
                            BuildFieldSimple( "isLTCG" ),
                            BuildFieldSimple( "isDataAligned" ),
                            BuildFieldSimple( "hasSecurityChecks" ),
                            BuildFieldSimple( "compilerName" ),
                            BuildFieldSimple( "hasAlloca" ),
                            BuildFieldSimple( "hasSetJump" ),
                            BuildFieldSimple( "hasLongJump" ),
                            BuildFieldSimple( "hasInlAsm" ),
                            BuildFieldSimple( "hasEH" ),
                            BuildFieldSimple( "hasSEH" ),
                            BuildFieldSimple( "hasEHa" ),
                            BuildFieldSimple( "isNaked" ),
                            BuildFieldSimple( "isAggregated" ),
                            BuildFieldSimple( "isSplitted" ),
                            BuildFieldSymbol( "container" ),
                            BuildFieldSimple( "inlSpec" ),
                            BuildFieldSimple( "noStackOrdering" ),
                            BuildFieldSymbol( "virtualBaseTableType" ),
                            BuildFieldSimple( "hasManagedCode" ),
                            BuildFieldSimple( "isHotpatchable" ),
                            BuildFieldSimple( "isCVTCIL" ),
                            BuildFieldSimple( "isMSILNetmodule" ),
                            BuildFieldSimple( "isCTypes" ),
                            BuildFieldSimple( "isStripped" ),
                          )

_LINE_FIELD_BUILDER     = ( BuildFieldSymbol( "compiland" ),
                            BuildFieldSrcFile( "sourceFile" ),
                            BuildFieldSimple( "lineNumber" ),
                            BuildFieldSimple( "lineNumberEnd" ),
                            BuildFieldSimple( "columnNumber" ),
                            BuildFieldSimple( "columnNumberEnd" ),
                            BuildFieldSimple( "addressSection" ),
                            BuildFieldHex( "addressOffset" ),
                            BuildFieldHex( "relativeVirtualAddress" ),
                            BuildFieldHex( "virtualAddress" ),
                            BuildFieldSimple( "length" ),
                            BuildFieldSimple( "statement" ),
                          )

_SRCFILE_FIELD_BUILDER  = ( BuildFieldSimple( "uniqueId" ),
                            BuildFieldSimple( "fileName" ),
                            BuildFieldEnumSymbols( "compilands" ),
                          )

# ------------------------------------------------------------------------------

class TreeItemBase(QtGui.QTreeWidgetItem):
    def __init__(self, parent):
        super(TreeItemBase, self).__init__(parent)

        font = self.font(0)
        font.setBold(True)
        self.setFont(0, font)

class TreeItemSymRva(TreeItemBase):
    def __init__(self, parent, rva):
        super(TreeItemSymRva, self).__init__(parent)
        self._rva = rva

    def getRva(self):
        return self._rva

class TreeItemExpandBase(TreeItemBase):
    def __init__(self, parent, val, builder, name):
        super(TreeItemExpandBase, self).__init__(parent)
        self._val = val
        self._builder = builder

        self._builded = False
        self.setChildIndicatorPolicy(QtGui.QTreeWidgetItem.ShowIndicator)

        self.setText(0, name)

    def onExpand(self):
        if self._builded:
            return
        self._builded = True

        for builder in self._builder:
            builder.do(self._val, self)

    def getVal(self):
        return self._val

class TreeItemSym(TreeItemExpandBase):
    def __init__(self, parent, sym):
        super(TreeItemSym, self).__init__( parent,
                                           sym,
                                           _SYMBOL_FIELD_BUILDER,
                                           "IDiaSymbol" )

    def getDiaSymbol(self):
        return self.getVal()

class TreeItemLine(TreeItemExpandBase):
    def __init__(self, parent, line):
        super(TreeItemLine, self).__init__( parent,
                                            line,
                                            _LINE_FIELD_BUILDER,
                                            "IDiaLineNumber" )

    def getDiaLineNumber(self):
        return self.getVal()

class TreeItemSrcFile(TreeItemExpandBase):
    def __init__(self, parent, srcfile):
        super(TreeItemSrcFile, self).__init__( parent,
                                               srcfile,
                                               _SRCFILE_FIELD_BUILDER,
                                               "IDiaSourceFile" )

    def getDiaSourceFile(self):
        return self.getVal()

# ------------------------------------------------------------------------------

class EnableMenuActions:
    def __init__(self, findChildren, findSourceLines):
        self.findChildren = findChildren
        self.findSourceLines = findSourceLines

class TreeBase(QtGui.QTreeWidget):
    def __init__(self, enableMenuActions):
        super(TreeBase, self).__init__()

        self.header().setResizeMode( QtGui.QHeaderView.Interactive )
        self.setHeaderLabels( ("", "") )

        self._enableMenuActions = enableMenuActions
        self.itemSelectionChanged.connect(self.onItemSelectionChanged)

        self.itemExpanded.connect(self.onItemExpanded)

    def onItemSelectionChanged(self):
        items = self.selectedItems()
        if 1 == len(items):
            self._enableMenuActions.findChildren( hasattr(items[0], "getDiaSymbol") )
            self._enableMenuActions.findSourceLines( hasattr(items[0], "getRva") )
        else:
            self._enableMenuActions.findChildren(False)
            self._enableMenuActions.findSourceLines(False)

    def onItemExpanded(self, item):
        if hasattr(item, "getVal"):
            item.onExpand()


class TreeSymbols(TreeBase):
    def __init__(self, enableMenuActions, collSyms, parentForProgress):
        super(TreeSymbols, self).__init__(enableMenuActions)

        progressDialog = QtGui.QProgressDialog( "Build",
                                                "Stop",
                                                0,
                                                len(collSyms) )
        progressDialog.setValue(0)

        for index in range( len(collSyms) ):
            if progressDialog.wasCanceled():
                break

            progressDialog.setValue(progressDialog.value() + 1)
            item = TreeItemSym(self, collSyms[index])
            self.addTopLevelItem(item)

class TreeLines(TreeBase):
    def __init__(self, enableMenuActions, collLines):
        super(TreeLines, self).__init__(enableMenuActions)
        for line in collLines:
            self.addTopLevelItem( TreeItemLine(self, line) )

# ------------------------------------------------------------------------------

class FindChildrenDialog(QtGui.QDialog):
    def __init__(self, parent):
        super(FindChildrenDialog, self).__init__(parent)

        self.editName = QtGui.QLineEdit()

        mapNameSearchOptions = diabind.NameSearchOptions.values
        self.lstNameSearchOptions = [ QtGui.QCheckBox( str(mapNameSearchOptions[i]) ) 
                                      for i in mapNameSearchOptions if i ]

        self.comboTag = QtGui.QComboBox()
        dictSymtag = diabind.SymTag.values
        for tag in dictSymtag:
            self.comboTag.addItem(str(dictSymtag[tag]))

        buttonBox = QtGui.QDialogButtonBox(QtGui.QDialogButtonBox.Ok | QtGui.QDialogButtonBox.Cancel)
        buttonBox.accepted.connect(self.accept)
        buttonBox.rejected.connect(self.reject)

        mainLayout = QtGui.QVBoxLayout()
        mainLayout.addWidget(QtGui.QLabel("Name:"))
        mainLayout.addWidget(self.editName)
        mainLayout.addWidget(self.comboTag)
        for checkBox in self.lstNameSearchOptions:
            mainLayout.addWidget(checkBox)
        mainLayout.addWidget(buttonBox)
        self.setLayout(mainLayout)

        self.setWindowTitle("IDiaSymbol::findChildren")

# ------------------------------------------------------------------------------

class MainWindow(QtGui.QMainWindow):
    def __init__(self, fileName):
        super(MainWindow, self).__init__()

        menu = self.menuBar().addMenu("&File")
        menu.addAction( QtGui.QAction( "&Open...", 
                                       self, 
                                       shortcut = "Ctrl+O",
                                       triggered = self.openFile ) )
        menu.addAction( QtGui.QAction( "&Exit...", 
                                       self, 
                                       shortcut = "Ctrl+E",
                                       triggered = self.close ) )

        menu = self.menuBar().addMenu("&Options")
        menu.addAction( QtGui.QAction( "&Symbol search path", 
                                       self, 
                                       shortcut = "Ctrl+S",
                                       triggered = self.setSearchPath ) )

        menu = self.menuBar().addMenu("&IDiaSymbol")
        self.actFindChildren = QtGui.QAction( "&Find children...", 
                                              self, 
                                              shortcut = "Ctrl+F",
                                              triggered = self.findChildren )
        menu.addAction( self.actFindChildren )

        menu = self.menuBar().addMenu("&relativeVirtualAddress")
        self.actFindSourceLines = QtGui.QAction( "&Find soure lines...", 
                                                 self, 
                                                 shortcut = "Ctrl+L",
                                                 triggered = self.findSourceLines )
        menu.addAction( self.actFindSourceLines )

        self.enableMenuActions = EnableMenuActions( self.actFindChildren.setEnabled,
                                                    self.actFindSourceLines.setEnabled )

        self.tabWidget = QtGui.QTabWidget()
        self.setCentralWidget(self.tabWidget)
        self.tabWidget.currentChanged.connect(self.onTabChanged)

        self.searchPath = ""

        self.resize(800, 600)

        self.reset("Ready. diabind " + diabind.version)

        if not fileName:
            self.openFile()
        else:
            self.loadFile(fileName)

    def onTabChanged(self, index):
        if -1 != index:
            self.tabWidget.widget(index).onItemSelectionChanged()

    def findChildren(self):
        sym = self.tabWidget.currentWidget().selectedItems()[0].getDiaSymbol()

        dialog = FindChildrenDialog(self)
        if dialog.exec_() != QtGui.QDialog.Accepted:
            return

        compareFlags = 0
        for checkBox in dialog.lstNameSearchOptions:
            if checkBox.isChecked():
                compareFlags = compareFlags | getattr(diabind.NameSearchOptions, checkBox.text())

        symTag = getattr(diabind.SymTag, dialog.comboTag.currentText())

        name = dialog.editName.text()

        collSyms = sym.findChildren(symTag, name, compareFlags)

        tabDesc =  "symtag=" + str(symTag)
        tabDesc += ", name=" + name
        tabDesc += ", compareFlags=" + hex(compareFlags)
        tabDesc += ". " + str(len(collSyms)) + " found"
        self.statusBar().showMessage(tabDesc)

        if not len(collSyms):
            QtGui.QMessageBox.information( self,
                                           "IDiaSymbol::findChildren",
                                           "Nothing found" )
            return

        self.addTab( TreeSymbols(self.enableMenuActions, collSyms, self),
                     str(symTag) + "|" + name + "|" + hex(compareFlags) )

    def findSourceLines(self):
        length, ok = QtGui.QInputDialog.getInteger(self, "IDiaSession::findLinesByRVA", "Length:")
        if not ok:
            return
        rva = self.tabWidget.currentWidget().selectedItems()[0].getRva()

        collLines = self.symSession.findLinesByRVA(rva, length)
        self.addTab( TreeLines(self.enableMenuActions, collLines),
                     hex(rva) + "+" + str(length) )

    def addTab(self, newTab, name):
        self.tabWidget.addTab(newTab, name)
        self.tabWidget.setCurrentWidget(newTab)

    def setSearchPath(self):
        searchPath, ok = QtGui.QInputDialog.getText( self, 
                                                     "Symbol search path",
                                                     "Alternate path to search for debug data:",
                                                     QtGui.QLineEdit.Normal,
                                                     self.searchPath )

        if ok:
            self.searchPath = searchPath

    def openFile(self):
        fileName = QtGui.QFileDialog.getOpenFileName( self,
                                                      "Open pdb\pe file...",
                                                      QtCore.QDir.currentPath(),
                                                      "Program database(*.pdb);;All files (*.*)")[0]
        if not fileName:
            return

        self.loadFile(fileName)

    def loadFile(self, fileName):
        self.reset("Ready. diabind " + diabind.version)

        self.symSession = None
        self.symGlobal = None

        try:
            if fileName[-4:].lower() == ".pdb":
                self.symSession = diabind.load_pdb(fileName)
            else:
                self.symSession = diabind.load_exe(fileName, self.searchPath)
        except diabind.DiaError as err:
            errDesc = str(err)
            self.reset(errDesc)
            QtGui.QMessageBox.critical( self,
                                        "Load debug information failed",
                                        errDesc )
            return

        self.addTab( TreeSymbols(self.enableMenuActions, (self.symSession.get_globalScope(), ), self),
                     "globalScope" )

        self.statusBar().showMessage(fileName + " loaded")
        self.setWindowTitle("diaview: " + fileName)

    def reset(self, status):
        self.actFindChildren.setEnabled(False)
        self.actFindSourceLines.setEnabled(False)
        self.tabWidget.clear()
        self.statusBar().showMessage(status)
        self.setWindowTitle("diaview")

# ------------------------------------------------------------------------------

if __name__ == '__main__':
    import sys
    app = QtGui.QApplication(sys.argv)
    mainWindow = MainWindow(sys.argv[1] if len(sys.argv) == 2 else "")
    mainWindow.show()
    app.exec_()
